# src/api/endpoints.py
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
import numpy as np
import pandas as pd
from datetime import datetime
import asyncio
import logging

# from src.models.portfolio import *
from src.models.portfolio import Objective, Constraint, OptimizationAlgorithm
from src.algorithms.change_detection import MarketChangeDetector
from src.algorithms.nsga2 import DynamicNSGA2
from src.algorithms.sac import SACPortfolioOptimizer
from src.data.market_data import MarketDataProcessor
from src.utils.performance import PerformanceCalculator
from src.visualization.plots import ParetoFrontVisualizer

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="动态多目标投资组合优化系统", version="1.0.0")

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局变量
market_detector = MarketChangeDetector()
data_processor = MarketDataProcessor()
performance_calculator = PerformanceCalculator()
visualizer = ParetoFrontVisualizer()

# 请求模型
class OptimizationRequest(BaseModel):
    assets: List[str]
    objectives: Objective
    constraints: Constraint
    algorithm: OptimizationAlgorithm
    start_date: str
    end_date: str

class MarketChangeRequest(BaseModel):
    returns_matrix: List[List[float]]
    test_method: str = "ks"

class SACTrainingRequest(BaseModel):
    assets: List[str]
    start_date: str
    end_date: str
    episodes: int = 1000
    
class VisualizationRequest(BaseModel):
    pareto_set: List[Dict[str, Any]]
    user_preferences: Dict[str, float]

# 响应模型
class OptimizationResponse(BaseModel):
    pareto_set: List[Dict[str, Any]]
    algorithm: str
    execution_time: float
    generation: int

class MarketChangeResponse(BaseModel):
    is_change: bool
    p_value: float
    trigger_at: datetime
    method: str

@app.post("/changes/detect", response_model=MarketChangeResponse)
async def detect_market_change(request: MarketChangeRequest):
    """市场变化检测接口"""
    try:
        returns_matrix = np.array(request.returns_matrix)
        
        result = market_detector.detect_market_change(
            returns_matrix, 
            method=request.test_method
        )
        
        return MarketChangeResponse(
            is_change=result['is_change'],
            p_value=result['p_value'],
            trigger_at=result['trigger_at'],
            method=result['method']
        )
        
    except Exception as e:
        logger.error(f"市场变化检测失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/optimize/run", response_model=OptimizationResponse)
async def run_optimization(request: OptimizationRequest):
    """运行优化算法接口"""
    try:
        start_time = datetime.now()
        
        # 获取市场数据
        market_data = data_processor.prepare_portfolio_data(
            request.assets, 
            request.start_date, 
            request.end_date
        )
        
        # 检测市场变化
        returns_matrix = market_data['returns'].values
        change_result = market_detector.detect_market_change(returns_matrix)
        
        # 选择优化算法
        if request.algorithm == OptimizationAlgorithm.NSGA2:
            optimizer = DynamicNSGA2(population_size=100, generations=50)
            
            # 构建约束字典
            constraints = {
                'max_single_weight': request.constraints.max_single_weight,
                'min_cash_ratio': request.constraints.min_cash_ratio
            }
            
            # 运行优化
            population = optimizer.optimize(
                returns=market_data['expected_returns'],
                cov_matrix=market_data['cov_matrix'],
                esg_scores=market_data['esg_scores'],
                constraints=constraints,
                market_changed=change_result['is_change']
            )
            
            # 转换为输出格式
            pareto_set = []
            for individual in population:
                if individual.rank == 0:  # 只返回第一前沿
                    solution = {
                        'weights': {
                            symbol: float(weight) 
                            for symbol, weight in zip(request.assets, individual.weights)
                        },
                        'metrics': {
                            'expected_return': float(individual.objectives[0]),
                            'risk': float(-individual.objectives[1]),
                            'esg_score': float(individual.objectives[2])
                        }
                    }
                    pareto_set.append(solution)
            
            execution_time = (datetime.now() - start_time).total_seconds()
            
            return OptimizationResponse(
                pareto_set=pareto_set,
                algorithm=request.algorithm,
                execution_time=execution_time,
                generation=optimizer.generations
            )
            
        elif request.algorithm == OptimizationAlgorithm.SAC:
            # SAC算法实现
            state_dim = len(request.assets) + 6  # 权重 + 市场指标
            action_dim = len(request.assets)
            
            sac_optimizer = SACPortfolioOptimizer(
                state_dim=state_dim,
                action_dim=action_dim
            )
            
            # 简化的SAC训练过程
            current_weights = np.ones(len(request.assets)) / len(request.assets)
            market_indicators = data_processor.get_market_indicators(datetime.now())
            
            state = sac_optimizer.get_state(current_weights, market_indicators)
            action = sac_optimizer.get_action(state, deterministic=True)
            
            # 归一化权重
            new_weights = sac_optimizer.normalize_weights(current_weights + action * 0.1)
            
            # 计算指标
            portfolio_return = np.dot(new_weights, market_data['expected_returns'])
            portfolio_risk = np.sqrt(np.dot(new_weights.T, 
                                          np.dot(market_data['cov_matrix'], new_weights)))
            esg_score = np.dot(new_weights, market_data['esg_scores'])
            
            solution = {
                'weights': {
                    symbol: float(weight) 
                    for symbol, weight in zip(request.assets, new_weights)
                },
                'metrics': {
                    'expected_return': float(portfolio_return),
                    'risk': float(portfolio_risk),
                    'esg_score': float(esg_score)
                }
            }
            
            execution_time = (datetime.now() - start_time).total_seconds()
            
            return OptimizationResponse(
                pareto_set=[solution],
                algorithm=request.algorithm,
                execution_time=execution_time,
                generation=1
            )
            
        else:
            raise HTTPException(status_code=400, detail=f"不支持的算法: {request.algorithm}")
            
    except Exception as e:
        logger.error(f"优化失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/train/sac")
async def train_sac(request: SACTrainingRequest, background_tasks: BackgroundTasks):
    """SAC训练接口"""
    try:
        # 在后台任务中训练SAC
        background_tasks.add_task(
            train_sac_model,
            request.assets,
            request.start_date,
            request.end_date,
            request.episodes
        )
        
        return {"message": "SAC训练已开始", "status": "started"}
        
    except Exception as e:
        logger.error(f"SAC训练启动失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

async def train_sac_model(assets: List[str], start_date: str, 
                         end_date: str, episodes: int):
    """SAC模型训练任务"""
    try:
        # 获取数据
        market_data = data_processor.prepare_portfolio_data(assets, start_date, end_date)
        
        # 初始化SAC
        state_dim = len(assets) + 6
        action_dim = len(assets)
        
        sac_optimizer = SACPortfolioOptimizer(
            state_dim=state_dim,
            action_dim=action_dim
        )
        
        # 训练循环
        current_weights = np.ones(len(assets)) / len(assets)
        returns_history = market_data['returns'].values
        
        for episode in range(episodes):
            # 随机选择一个时间点
            t = np.random.randint(30, len(returns_history) - 1)
            
            # 获取当前状态
            market_indicators = {
                'vix': np.random.normal(20, 5),
                'market_return': returns_history[t].mean(),
                'market_volatility': returns_history[max(0, t-30):t].std(),
                'rsi': 50 + 30 * np.sin(t / 10),
                'macd': np.random.normal(0, 0.5),
                'volume_ratio': np.random.normal(1.0, 0.2)
            }
            
            state = sac_optimizer.get_state(current_weights, market_indicators)
            
            # 获取动作
            action = sac_optimizer.get_action(state)
            
            # 计算新权重
            new_weights = sac_optimizer.normalize_weights(current_weights + action * 0.1)
            
            # 计算奖励
            next_returns = returns_history[t + 1]
            reward = sac_optimizer.calculate_reward(
                current_weights, new_weights, next_returns
            )
            
            # 获取下一状态
            next_market_indicators = {
                'vix': np.random.normal(20, 5),
                'market_return': returns_history[t + 1].mean(),
                'market_volatility': returns_history[max(0, t-29):t+1].std(),
                'rsi': 50 + 30 * np.sin((t + 1) / 10),
                'macd': np.random.normal(0, 0.5),
                'volume_ratio': np.random.normal(1.0, 0.2)
            }
            
            next_state = sac_optimizer.get_state(new_weights, next_market_indicators)
            
            # 存储经验
            done = episode == episodes - 1
            sac_optimizer.store_transition(state, action, reward, next_state, done)
            
            # 更新网络
            if episode > 100:  # 预热
                sac_optimizer.update_networks()
            
            # 更新权重
            current_weights = new_weights
            
            # 每100轮打印一次进度
            if episode % 100 == 0:
                logger.info(f"SAC训练进度: {episode}/{episodes}, 奖励: {reward:.4f}")
        
        # 保存模型
        model_path = f"models/sac_model_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pt"
        sac_optimizer.save_model(model_path)
        
        logger.info(f"SAC训练完成，模型已保存至: {model_path}")
        
    except Exception as e:
        logger.error(f"SAC训练失败: {e}")

@app.post("/visualize/pareto")
async def visualize_pareto_front(request: VisualizationRequest):
    """Pareto前沿可视化接口"""
    try:
        # 生成可视化数据
        viz_data = visualizer.create_pareto_visualization(
            request.pareto_set,
            request.user_preferences
        )
        
        return {"visualization_data": viz_data}
        
    except Exception as e:
        logger.error(f"可视化失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/alerts/slack")
async def send_slack_alert(alert_type: str, details: Dict[str, Any]):
    """Slack告警接口"""
    try:
        # 模拟发送Slack消息
        message_id = f"msg_{datetime.now().timestamp()}"
        
        logger.info(f"发送Slack告警: {alert_type}, 详情: {details}")
        
        return {"message_id": message_id}
        
    except Exception as e:
        logger.error(f"Slack告警发送失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """健康检查接口"""
    return {"status": "healthy", "timestamp": datetime.now()}

@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "🚀 动态多目标投资组合优化系统",
        "version": "1.0.0",
        "endpoints": [
            "/changes/detect",
            "/optimize/run",
            "/train/sac",
            "/visualize/pareto",
            "/alerts/slack"
        ]
    }
